import os
import shutil
from concurrent.futures import ThreadPoolExecutor, as_completed

# 源路径
source_base = r"E:\zhihao\2025文档\25年10月\ai_训练\ai项目_备份\yangzhihao\ai_project\public_security\traffic_accient\train\traffic_accidents_model\datasets\valid"
source_images = os.path.join(source_base, "images")
source_labels = os.path.join(source_base, "labels")

# 目标路径
lyh_base = r"E:\zhihao\2025文档\25年10月\ai_训练\数据集\重清洗\lyh\valid"
hx_base = r"E:\zhihao\2025文档\25年10月\ai_训练\数据集\重清洗\hx\valid"

lyh_images = os.path.join(lyh_base, "images")
lyh_labels = os.path.join(lyh_base, "labels")
hx_images = os.path.join(hx_base, "images")
hx_labels = os.path.join(hx_base, "labels")

# 创建目标文件夹
for folder in [lyh_images, lyh_labels, hx_images, hx_labels]:
    os.makedirs(folder, exist_ok=True)

print("✅ 已创建所有目标文件夹")

# 获取所有图片文件(排除包含 "yx" 的)
print("\n扫描文件...")
image_files = []
for filename in os.listdir(source_images):
    if os.path.isfile(os.path.join(source_images, filename)):
        if "yx" not in filename.lower():
            image_files.append(filename)

print(f"找到 {len(image_files)} 个不包含 'yx' 的图片文件")

# 平分文件
mid = len(image_files) // 2
lyh_files = image_files[:mid]
hx_files = image_files[mid:]

print(f"  - lyh: {len(lyh_files)} 个文件")
print(f"  - hx: {len(hx_files)} 个文件")


def copy_file_pair(filename, src_img_dir, src_lbl_dir, dst_img_dir, dst_lbl_dir):
    """复制一对文件(图片和对应的标签)"""
    try:
        # 复制图片
        src_img = os.path.join(src_img_dir, filename)
        dst_img = os.path.join(dst_img_dir, filename)
        shutil.copy2(src_img, dst_img)

        # 查找并复制对应的标签文件
        base_name = os.path.splitext(filename)[0]
        label_copied = False

        # 尝试常见的标签扩展名
        for ext in ['.txt', '.xml', '.json']:
            label_file = base_name + ext
            src_lbl = os.path.join(src_lbl_dir, label_file)
            if os.path.exists(src_lbl):
                dst_lbl = os.path.join(dst_lbl_dir, label_file)
                shutil.copy2(src_lbl, dst_lbl)
                label_copied = True
                break

        return (True, label_copied)
    except Exception as e:
        return (False, False)


# 使用多线程加速复制
print("\n开始复制文件...")

# 复制 lyh 的文件
print("复制 lyh 的文件...")
lyh_success = 0
lyh_labels_found = 0

with ThreadPoolExecutor(max_workers=8) as executor:
    futures = [
        executor.submit(copy_file_pair, f, source_images, source_labels, lyh_images, lyh_labels)
        for f in lyh_files
    ]

    for i, future in enumerate(as_completed(futures), 1):
        img_ok, lbl_ok = future.result()
        if img_ok:
            lyh_success += 1
        if lbl_ok:
            lyh_labels_found += 1

        if i % 100 == 0 or i == len(lyh_files):
            print(f"  进度: {i}/{len(lyh_files)}")

# 复制 hx 的文件
print("\n复制 hx 的文件...")
hx_success = 0
hx_labels_found = 0

with ThreadPoolExecutor(max_workers=8) as executor:
    futures = [
        executor.submit(copy_file_pair, f, source_images, source_labels, hx_images, hx_labels)
        for f in hx_files
    ]

    for i, future in enumerate(as_completed(futures), 1):
        img_ok, lbl_ok = future.result()
        if img_ok:
            hx_success += 1
        if lbl_ok:
            hx_labels_found += 1

        if i % 100 == 0 or i == len(hx_files):
            print(f"  进度: {i}/{len(hx_files)}")

# 总结
print("\n" + "=" * 60)
print("✅ 任务完成!")
print(f"\nlyh 统计:")
print(f"  - 图片: {lyh_success} 个")
print(f"  - 标签: {lyh_labels_found} 个")
print(f"\nhx 统计:")
print(f"  - 图片: {hx_success} 个")
print(f"  - 标签: {hx_labels_found} 个")
print(f"\n总计:")
print(f"  - 图片: {lyh_success + hx_success} 个")
print(f"  - 标签: {lyh_labels_found + hx_labels_found} 个")
print("=" * 60)